import pandas as pd
import os
from datetime import datetime, timedelta
import plotly.graph_objects as go
from plotly.subplots import make_subplots


# --- 数据加载和处理函数 (与之前版本相同，无需修改) ---
def load_data_in_range(start_time_str: str, end_time_str: str, base_path: str, symbol: str,
                       timeframe: str) -> pd.DataFrame:
    """从指定的CSV文件中加载时间范围内的数据，并动态计算Volume Delta。"""
    start_date = pd.to_datetime(start_time_str)
    end_date = pd.to_datetime(end_time_str)
    data_folder_path = os.path.join(base_path, symbol, timeframe)
    if not os.path.isdir(data_folder_path):
        print(f"错误: 文件夹不存在 -> {data_folder_path}")
        return pd.DataFrame()

    date_range = pd.date_range(start=start_date.date(), end=end_date.date(), freq='D').strftime('%Y-%m').unique()
    if len(date_range) == 0: date_range = [start_date.strftime('%Y-%m')]
    all_dfs = []
    print(f"将在以下月份文件中查找数据: {list(date_range)}")
    for year_month in date_range:
        file_path = os.path.join(data_folder_path, f"{symbol.upper()}-{timeframe}-{year_month}.csv")
        if os.path.exists(file_path):
            print(f"正在读取文件: {file_path}")
            try:
                df = pd.read_csv(file_path)
                required_cols = ['open_time', 'open', 'high', 'low', 'close', 'volume', 'taker_buy_volume']
                if not all(col in df.columns for col in required_cols):
                    print(f"错误: 文件 {file_path} 缺少必要的列。")
                    continue

                for col in ['volume', 'taker_buy_volume']:
                    df[col] = pd.to_numeric(df[col], errors='coerce')

                df['volume_delta'] = (2 * df['taker_buy_volume']) - df['volume']
                all_dfs.append(df)
            except Exception as e:
                print(f"读取或处理文件 {file_path} 时出错: {e}")
        else:
            print(f"警告: 未找到文件 {file_path}")

    if not all_dfs:
        print("错误: 在指定时间范围内未找到任何数据文件。")
        return pd.DataFrame()

    combined_df = pd.concat(all_dfs, ignore_index=True)
    combined_df.drop_duplicates(subset=['open_time'], inplace=True)
    return combined_df


def process_and_filter_data(df: pd.DataFrame, start_time_str: str, end_time_str: str) -> pd.DataFrame:
    if df.empty:
        return df
    df['open_time'] = pd.to_datetime(df['open_time'], unit='ms')
    df.set_index('open_time', inplace=True)
    numeric_cols = ['open', 'high', 'low', 'close', 'volume', 'volume_delta']
    df[numeric_cols] = df[numeric_cols].apply(pd.to_numeric, errors='coerce')
    df.dropna(subset=numeric_cols, inplace=True)
    df.sort_index(inplace=True)
    df['cvd'] = df['volume_delta'].cumsum()
    # 使用 pd.to_datetime 转换字符串时间以确保时区一致性
    start_ts = pd.to_datetime(start_time_str)
    end_ts = pd.to_datetime(end_time_str)
    df_filtered = df.loc[start_ts:end_ts]
    print(f"数据处理完成。筛选出 {len(df_filtered)} 条K线。")
    return df_filtered


# --- 新增：枢轴点计算函数 ---
def calculate_pivot_points(df_24h: pd.DataFrame) -> dict:
    """
    根据过去24小时的数据计算枢轴点。
    使用最后一条K线的收盘价作为'Close'。
    """
    if df_24h.empty:
        print("用于计算枢轴点的数据为空，无法计算。")
        return {}

    high_24h = df_24h['high'].max()
    low_24h = df_24h['low'].min()
    # 获取时间序列中最后一根K线的收盘价
    close_24h = df_24h['close'].iloc[-1]

    print(f"\n--- 枢轴点计算依据 ---")
    print(f"过去24小时最高价 (High): {high_24h:.2f}")
    print(f"过去24小时最低价 (Low): {low_24h:.2f}")
    print(f"过去24小时收盘价 (Close): {close_24h:.2f}")

    pp = (high_24h + low_24h + close_24h) / 3
    s1 = (pp * 2) - high_24h
    r1 = (pp * 2) - low_24h
    s2 = pp - (high_24h - low_24h)
    r2 = pp + (high_24h - low_24h)
    s3 = low_24h - 2 * (high_24h - pp)
    r3 = high_24h + 2 * (pp - low_24h)

    pivots = {
        'R3': r3, 'R2': r2, 'R1': r1,
        'PP': pp,
        'S1': s1, 'S2': s2, 'S3': s3
    }

    print("\n计算出的枢轴点位:")
    for name, value in pivots.items():
        print(f"  {name}: {value:.2f}")

    return pivots


# --- 绘图函数 (已更新以接收并绘制枢轴点) ---
def plot_with_plotly_legacy_compatible(df: pd.DataFrame, chart_title: str, pivots: dict = None):
    """
    使用旧版 Plotly 兼容的方式绘制图表 (使用 hovertext)。
    新增功能：在K线图上绘制枢轴点水平线。
    """
    if df.empty:
        print("数据为空，无法绘制图表。")
        return

    hover_texts = []
    for index, row in df.iterrows():
        text = (f"<b>Time</b>: {index.strftime('%Y-%m-%d %H:%M')}<br>"
                f"<b>Open</b>: {row['open']:.2f}<br>"
                f"<b>High</b>: {row['high']:.2f}<br>"
                f"<b>Low</b>: {row['low']:.2f}<br>"
                f"<b>Close</b>: {row['close']:.2f}<br>"
                f"<b>Volume</b>: {row['volume']:,.0f}<br>"
                f"<b>Vol Delta</b>: {row['volume_delta']:.2f}<br>"
                f"<b>CVD</b>: {row['cvd']:,.2f}")
        hover_texts.append(text)

    fig = make_subplots(
        rows=4, cols=1,
        shared_xaxes=True,
        vertical_spacing=0.02,
        row_heights=[0.6, 0.1, 0.1, 0.2]
    )

    # 1. K线图
    fig.add_trace(go.Candlestick(x=df.index,
                                 open=df['open'], high=df['high'],
                                 low=df['low'], close=df['close'],
                                 name='OHLC',
                                 hoverinfo='text',
                                 hovertext=hover_texts),
                  row=1, col=1)

    # --- 新增：添加枢轴点水平线 ---
    if pivots:
        colors = {'R3': 'darkred', 'R2': 'red', 'R1': 'lightcoral', 'PP': 'blue',
                  'S1': 'lightgreen', 'S2': 'green', 'S3': 'darkgreen'}
        line_styles = {'PP': 'solid', 'R1': 'dash', 'S1': 'dash', 'R2': 'dot', 'S2': 'dot', 'R3': 'dashdot',
                       'S3': 'dashdot'}

        for name, value in pivots.items():
            fig.add_hline(
                y=value,
                line_dash=line_styles.get(name, "dash"),
                annotation_text=f"{name}: {value:.2f}",
                annotation_position="bottom right",
                annotation_font_size=12,
                annotation_font_color=colors.get(name, 'grey'),
                line_color=colors.get(name, 'grey'),
                row=1, col=1
            )

    # 2. 成交量图
    fig.add_trace(go.Bar(x=df.index, y=df['volume'], name='Volume',
                         marker_color='lightblue',
                         hoverinfo='text',
                         hovertext=hover_texts),
                  row=2, col=1)

    # 3. Volume Delta 图
    delta_colors = ['green' if v >= 0 else 'red' for v in df['volume_delta']]
    fig.add_trace(go.Bar(x=df.index, y=df['volume_delta'], name='Volume Delta',
                         marker_color=delta_colors,
                         hoverinfo='text',
                         hovertext=hover_texts),
                  row=3, col=1)

    # 4. CVD 图
    fig.add_trace(go.Scatter(x=df.index, y=df['cvd'], name='CVD',
                             line=dict(color='purple'),
                             hoverinfo='text',
                             hovertext=hover_texts),
                  row=4, col=1)

    # 更新布局
    start_str = df.index[0].strftime('%Y-%m-%d %H:%M')
    end_str = df.index[-1].strftime('%Y-%m-%d %H:%M')
    full_title = f'{chart_title} ({start_str} to {end_str}) with Pivot Points'

    fig.update_layout(
        title_text=full_title,
        xaxis_rangeslider_visible=False,
        showlegend=False,
        yaxis1_title='Price',
        yaxis2_title='Volume',
        yaxis3_title='Vol Delta',
        yaxis4_title='CVD',
        hovermode='x unified'
    )

    fig.show()


if __name__ == '__main__':
    # ==================== 用户配置区 ====================
    BASE_DATA_PATH = "F:/personal/binance_klines"
    SYMBOL = "SUIUSDT"
    TIMEFRAME = "1h"
    # 加载更多数据来测试性能
    START_TIME = "2025-08-05 00:00:00"
    END_TIME = "2025-08-07 00:00:00"
    # ====================================================

    print("--- K线图表程序启动 ---")

    # --- 新增：计算枢轴点的步骤 ---
    # 1. 定义用于计算枢轴点的时间范围（图表开始时间的前24小时）
    start_dt = pd.to_datetime(START_TIME)
    pivot_end_time_dt = start_dt
    pivot_start_time_dt = pivot_end_time_dt - timedelta(hours=24)

    pivot_start_str = pivot_start_time_dt.strftime('%Y-%m-%d %H:%M:%S')
    pivot_end_str = pivot_end_time_dt.strftime('%Y-%m-%d %H:%M:%S')

    # 2. 加载这24小时的数据
    print(f"\n--- 正在加载过去24小时数据以计算枢轴点 ({pivot_start_str} to {pivot_end_str}) ---")
    pivot_raw_data = load_data_in_range(pivot_start_str, pivot_end_str, BASE_DATA_PATH, SYMBOL, TIMEFRAME)
    pivot_df = process_and_filter_data(pivot_raw_data, pivot_start_str, pivot_end_str)

    # 3. 计算枢轴点
    pivot_points = {}
    if not pivot_df.empty:
        pivot_points = calculate_pivot_points(pivot_df)
    else:
        print("\n警告: 未能加载用于计算枢轴点的数据，图表中将不显示枢轴线。")
    # --- 枢轴点计算结束 ---

    print(f"\n--- 正在加载图表数据 ({START_TIME} to {END_TIME}) ---")
    raw_data_df = load_data_in_range(START_TIME, END_TIME, BASE_DATA_PATH, SYMBOL, TIMEFRAME)
    final_df = process_and_filter_data(raw_data_df, START_TIME, END_TIME)

    if not final_df.empty:
        chart_title_str = f"{SYMBOL} - {TIMEFRAME}"
        # 将计算出的枢轴点传入绘图函数
        plot_with_plotly_legacy_compatible(final_df, chart_title=chart_title_str, pivots=pivot_points)
        print("\n--- 图表生成完毕 ---")
    else:
        print("\n未能加载或处理任何图表数据，程序退出。请检查您的配置和文件路径。")